# src/template_service.py
import logging
import os
from typing import List, Optional

import aiofiles
import pandas as pd
from fastapi import UploadFile

from reportagentic.core.db import TemplateRepository, MetadataRepository
from reportagentic.core.llm import LLMClient
from reportagentic.core.vector import VectorClient


class TemplateService:
    def __init__(self, template_repo: TemplateRepository,
                 meta_repo: MetadataRepository,
                 llm_client: LLMClient,
                 vector_client: VectorClient,
                 upload_dir: str):
        self.template_repo = template_repo
        self.meta_repo = meta_repo
        self.llm_client = llm_client
        self.vector_client = vector_client
        self.upload_dir = upload_dir
        os.makedirs(self.upload_dir, exist_ok=True)

    async def _save_and_parse_template_file(self, file: UploadFile) -> tuple[str, list]:
        """
        保存文件并使用Pandas解析其头部。
        """
        file_path = os.path.join(self.upload_dir, file.filename)
        async with aiofiles.open(file_path, 'wb') as f:
            content = await file.read()
            await f.write(content)

        try:
            if file_path.endswith('.csv'):
                df = pd.read_csv(file_path)
            elif file_path.endswith('.xls'):
                df = pd.read_excel(file_path, engine='xlrd')
            elif file_path.endswith('.xlsx'):
                df = pd.read_excel(file_path, engine='openpyxl')
            else:
                raise ValueError(f'Unsupported file type: {file_path}')
            headers = df.columns.tolist()
            return file_path, headers
        except Exception as e:
            logging.error(f"Failed to parse file {file.filename}: {e}")
            raise ValueError(f"Could not parse the template file. Please ensure it is a valid Excel or CSV file.")

    async def register_template(
            self,
            template_name: str,
            description: str,
            template_file: UploadFile,
            example_descriptions: Optional[List[str]]
    ) -> dict:
        """
        注册模板，并调用LLM自动完成字段映射。
        """
        # 1. 解析模板文件，获取表头
        template_file_path, template_headers = await self._save_and_parse_template_file(template_file)
        logging.info(f"Template headers parsed: {template_headers}")

        # 2. 字段映射
        logging.info("Mapping template fields to database schema using LLM...")
        db_schema_string = self.meta_repo.get_schema_as_string()
        if not db_schema_string:
            raise ValueError("Database schema is not available. Cannot perform mapping.")
        mappings = self.llm_client.map_template_to_schema(template_headers, db_schema_string)
        if not mappings or len(mappings) != len(template_headers):
            raise ValueError("LLM failed to generate valid mappings for all template fields.")
        logging.info(f"Field mapping successful: {mappings}")

        # 3. 保存主模板信息，包括用例描述
        template_id = self.template_repo.save_template(template_name, description, example_descriptions,
                                                       template_file_path)
        logging.info(f"Registered base template '{template_name}' with ID {template_id}.")

        # 4. 保存成功的映射关系到数据库
        self.template_repo.save_mapped_parameters(template_id, mappings)
        logging.info(f"Saved {len(mappings)} field mappings for template ID {template_id}.")

        # 5. 为模板本身创建向量索引
        # 将所有文本信息合并，创建一个综合的向量
        all_text_parts = [f"报表模板: {template_name}", f"功能描述: {description}"]
        if example_descriptions:
            all_text_parts.append("用例: " + " | ".join(example_descriptions))

        text_to_embed = "\n".join(all_text_parts)
        vector = self.llm_client.get_embeddings([text_to_embed])[0]

        self.vector_client.upsert_vectors(
            collection_name="template_vectors",
            ids=[template_id],
            vectors=[vector],
            payloads=[{"name": template_name, "description": description}]
        )
        logging.info(f"Indexed vector for template ID {template_id}.")
        return {"template_id": template_id, "field_mappings": mappings}
